# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import re
import os
import pprint
from enum import Enum

from coin.base.datetime_util import to_datetime
from coin.strategy.tool.log_v1.util import find_log_pos_by_timestamp
from coin.strategy.tool.log_v2.util import (extract_datetime, parse_fill_dump, parse_fill_log)


def diff_or_none(dt2, dt1, modifier=None):
  if dt2 is None or dt1 is None:
    return None
  modifier = modifier or (lambda x: x)
  return modifier(dt2) - modifier(dt1)


class StatusSnapshotParser(object):
  def __init__(self, line):
    pattern = '^\+[ ]+([a-zA-Z0-9_\-.]+)([ ]+(\S.*))?$'
    self._matcher = re.compile(pattern)

    self._snapshot = {}
    self._header_line_count = 0
    self._done = False
    self._parse_header(line)

  @property
  def done(self):
    return self._done

  @property
  def status_snapshot(self):
    return self._snapshot

  def _parse_header(self, line):
    assert not self._done
    if self._header_line_count == 0:
      assert line.startswith('+++ STATUS: ')
      self._snapshot['time'] = (datetime.datetime.strptime(line[12:], '%Y-%m-%d %H:%M:%S.%f'))
    self._header_line_count += 1

  def parse_line(self, line):
    assert not self._done
    if self._header_line_count <= 2:
      self._parse_header(line)
      return

    if not line.startswith('+'):
      raise ValueError()

    matched = self._matcher.match(line)
    if matched is None:
      assert line == '+++', line
      self._done = True
      return

    key = matched.group(1)
    value = matched.group(3)
    self._snapshot[key] = value


def get_snapshot_time(current_time, mark_time_margin=None):
  mark_time_margin = mark_time_margin or datetime.timedelta(minutes=10)
  return {
      'latest': current_time,
      'mark': (current_time - mark_time_margin).replace(hour=0, minute=0, second=0, microsecond=0),
      '4h': current_time - datetime.timedelta(hours=4),
      '24h': current_time - datetime.timedelta(hours=24)
  }


def to_float_or_none(x):
  return float(x) if x is not None else None


def get_snapshot_balance(snapshot, quote):
  margin_balance_str = snapshot['balance.%s' % quote]
  if (margin_balance_str is None or float(margin_balance_str) < 1e-6):
    return None

  try:
    return to_float_or_none(snapshot['balance.%s.total' % quote])
  except KeyError:
    return to_float_or_none(snapshot['balance.total'])


def get_snapshot_position(snapshot, product_str):
  try:
    return to_float_or_none(snapshot['position.%s' % product_str])
  except KeyError:
    return to_float_or_none(snapshot['product.%s.position' % product_str])


def list_balance_currencies(snapshot):
  balance_curs = set()
  for key in snapshot.keys():
    matched = re.match('^balance\.([A-Z]+)$', key)
    if matched:
      balance_curs.add(matched.group(1))
  return sorted(balance_curs)


def list_position_products(snapshot):
  products = set()
  for key in snapshot.keys():
    for pattern, group_idx in [('^position\.([0-9A-Za-z.-]+)$', 1),
                               ('^product\.([0-9A-Za-z.-]+)\.position$', 1)]:
      matched = re.match(pattern, key)
      if matched:
        products.add(matched.group(group_idx))
        break
  if not products:
    pprint.pprint(snapshot)
  return sorted(products)


def change_dict_key(dict_obj, key, new_key, ignore_keyerror=True):
  try:
    value = dict_obj[key]
    del dict_obj[key]
    dict_obj[new_key] = value
  except KeyError:
    if not ignore_keyerror:
      raise


class LogParser(object):
  class State(Enum):
    REGULAR = 0
    IN_STATUS = 1
    DONE = 10

  class ParseFinished(Exception):
    pass

  def __init__(self,
               log_path,
               time_end=None,
               time_begin=None,
               mark_time_margin=None,
               dump_balance_position=False,
               auto_retires_on_name_change=True):
    self._log_path = os.path.normpath(os.path.expanduser(log_path))
    self._time_end = time_end or datetime.datetime.now()
    self._time_begin = (time_begin or self._time_end - datetime.timedelta(hours=26))
    self._last_log_time = None

    self._snapshot_time = get_snapshot_time(time_end, mark_time_margin)
    self._snapshot_map = {}
    self._turnover_map = {}
    self._last_fill_map = {}

    self._dump_balance_position = dump_balance_position
    self._balance_dump = {}
    self._position_dump = {}

    self._state = LogParser.State.REGULAR
    self._snapshot_parser = None

    self._auto_retires_on_name_change = auto_retires_on_name_change
    self._current_strategy_name = None

    self._processing_time = None
    self._line_counter = 0

  def _parse_fill(self, line):
    for fill_parser in (parse_fill_dump, parse_fill_log):
      fill = fill_parser(line)
      if fill is not None:
        if fill.timestamp > self._time_end:
          raise LogParser.ParseFinished()
        return fill
    return None

  def _parse_regular(self, line):
    if line.startswith('+++ STATUS: '):
      try:
        self._snapshot_parser = StatusSnapshotParser(line)
        self._state = LogParser.State.IN_STATUS
      except Exception:  # Ignore error.
        pass
    else:
      fill = self._parse_fill(line)
      if fill is not None:
        self._on_fill(fill)
        return

      if self._line_counter % 1000 == 0:
        line_timestamp = extract_datetime(line)
        if line_timestamp is not None and line_timestamp > self._time_end:
          raise LogParser.ParseFinished()

  def _parse_status(self, line):
    assert self._snapshot_parser
    try:
      self._snapshot_parser.parse_line(line)
      if self._snapshot_parser.done:
        self._on_status_snapshot(self._snapshot_parser.status_snapshot)
        self._snapshot_parser = None
        self._state = LogParser.State.REGULAR
    except LogParser.ParseFinished:
      self._snapshot_parser = None
      raise
    except Exception:  # Ignore error.
      self._snapshot_parser = None
      self._state = LogParser.State.REGULAR

  def _on_status_snapshot(self, snapshot):
    if snapshot['time'] > self._time_end:
      raise LogParser.ParseFinished()

    name = snapshot.get('name', None)
    if self._current_strategy_name != name:
      self._on_strategy_changed(name)
      self._current_strategy_name = name

    self._current_strategy_name = name
    self._update_snapshot_map(name, snapshot)

  def _check_status_snapshot_sanity(self, snapshot):
    quote = list_balance_currencies(snapshot)[0]
    balance = get_snapshot_balance(snapshot, quote)
    if balance is None:
      return False

    product = list_position_products(snapshot)[-1]
    position = get_snapshot_position(snapshot, product)
    if position is None:
      return False

    return True

  def _on_strategy_changed(self, new_name):
    if (not self._auto_retires_on_name_change or self._current_strategy_name is None
        or new_name is None):
      return

    key = self._current_strategy_name
    new_key = 'RETIRED - %s' % self._current_strategy_name
    change_dict_key(self._snapshot_map, key, new_key)
    change_dict_key(self._turnover_map, key, new_key)
    change_dict_key(self._last_fill_map, key, new_key)
    change_dict_key(self._balance_dump, key, new_key)
    change_dict_key(self._position_dump, key, new_key)

  def _update_snapshot_map(self, name, snapshot):
    if name is None:
      return

    if not self._check_status_snapshot_sanity(snapshot):
      return

    if name not in self._snapshot_map:
      self._snapshot_map[name] = {k: None for k, _ in self._snapshot_time.items()}

    if name not in self._turnover_map:
      self._turnover_map[name] = {k: 0. for k, _ in self._snapshot_time.items()}

    snapshot_time = snapshot['time']
    for time_key, time in self._snapshot_time.items():
      if (snapshot_time <= time or self._snapshot_map[name][time_key] is None):
        self._snapshot_map[name][time_key] = snapshot

    if self._dump_balance_position:
      if name not in self._balance_dump:
        self._balance_dump[name] = []
        self._position_dump[name] = []

      quote = list_balance_currencies(snapshot)[0]
      product = list_position_products(snapshot)[-1]
      balance = get_snapshot_balance(snapshot, quote)
      if balance is not None:
        self._balance_dump[name].append({'time': snapshot_time, 'value': balance})
      position = get_snapshot_position(snapshot, product)
      if position is not None:
        self._position_dump[name].append({'time': snapshot_time, 'value': position})

  def _on_fill(self, fill):
    if fill.timestamp > self._time_end:
      raise LogParser.ParseFinished()

    name = self._current_strategy_name
    if name is None:
      return

    self._last_fill_map[name] = fill

    if name not in self._turnover_map:
      self._turnover_map[name] = {k: 0. for k, _ in self._snapshot_time.items()}

    for time_key, time in self._snapshot_time.items():
      if fill.timestamp > time:
        self._turnover_map[name][time_key] += fill.notional_amt

  def parse_log_line(self, line):
    self._last_log_time = extract_datetime(line) or self._last_log_time
    try:
      if self._state == LogParser.State.REGULAR:
        self._parse_regular(line)
      elif self._state == LogParser.State.IN_STATUS:
        self._parse_status(line)
      else:
        raise LogParser.ParseFinished()

    except LogParser.ParseFinished:
      self._state = LogParser.State.DONE

    finally:
      self._line_counter += 1

  @property
  def done(self):
    return self._state == LogParser.State.DONE

  @property
  def status_snapshots(self):
    return self._snapshot_map

  @property
  def latest_strategy_name(self):
    return self._current_strategy_name

  def get_strategy_summary(self, name):
    snapshots = self._snapshot_map[name]
    turnovers = self._turnover_map[name]

    sn_latest = snapshots['latest']
    sn_mark = snapshots['mark']
    sn_4h = snapshots['4h']
    sn_24h = snapshots['24h']

    quote = list_balance_currencies(sn_latest)[0]
    product = list_position_products(sn_latest)[-1]
    get_balance = lambda x: get_snapshot_balance(x, quote)

    value_latest = get_balance(sn_latest)
    position_latest = get_snapshot_position(sn_latest, product)

    summary = {
        'name': name,
        'group': sn_latest.get('group', None),
        'quote': quote,
        'product': product,
        'account_value': value_latest,
        'mark_value': get_balance(sn_mark),
        'pnl_mark': diff_or_none(value_latest, get_balance(sn_mark), float),
        'pnl_4h': diff_or_none(value_latest, get_balance(sn_4h), float),
        'pnl_24h': diff_or_none(value_latest, get_balance(sn_24h), float),
        'turnover_mark': turnovers['mark'],
        'turnover_4h': turnovers['4h'],
        'turnover_24h': turnovers['24h'],
        'position': position_latest,
        'last_fill': self._last_fill_map.get(name, None),
        'mark_time': sn_mark['time'],
        'balance_dump': None,
        'position_dump': None,
        'processing_time': self._processing_time,
        'last_log': self._last_log_time
    }

    if self._dump_balance_position:
      summary['balance_dump'] = self._balance_dump[name]
      summary['position_dump'] = self._position_dump[name]
    return summary

  def get_summary(self):
    return [self.get_strategy_summary(name) for name in self._snapshot_map.keys()]

  def run(self, use_pos_cache=False):
    begin_dt = datetime.datetime.now()
    try:
      pos = find_log_pos_by_timestamp(self._log_path, self._time_begin, use_cache=use_pos_cache)
    except FileNotFoundError:
      return self.get_summary()

    with open(self._log_path, 'r', errors='ignore') as log_file:
      log_file.seek(pos)
      for line in log_file.readlines():
        self.parse_log_line(line.strip())
        if self._state == LogParser.State.DONE:
          break

    end_dt = datetime.datetime.now()
    self._processing_time = end_dt - begin_dt
    return self.get_summary()
